# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/28 18:50 
@Author  : ZhangShenao 
@File    : sql_generator.py
@Desc    : 利用OpenAI的工具调用,生成SQL

执行步骤:
1. 获取数据库的schema信息,包括表名及其包含的列名
2. 用户输入自然语言的提问query
3. 调用LLM,利用tool calling功能,将schema信息作为工具参数传入,LLM会基于用户提问,生成结构化的SQL
4. 执行LLM生成的SQL,获取最终结果

最终效果:
用户使用自然语言提问,即可以生成标准的结构化的SQL,执行后获取结果
这是一种编程范式的升级

"""
import json
import os
import sqlite3
from sqlite3 import Connection
from typing import List

import dotenv
from openai import OpenAI


def connect_db() -> Connection:
    """创建sqlite数据库连接"""
    conn = sqlite3.connect("../data/chinook.db")
    print("connect to sqlite success!")
    return conn


def get_table_names(conn: Connection) -> List[str]:
    """返回一个包含所有表名的列表"""
    # 执行SQL查询,获取数据库中所有表的名字
    tables = conn.execute("SELECT name FROM sqlite_master WHERE type='table';")
    # 遍历查询table,返回table_name列表
    return [table[0] for table in tables]


def get_column_names(conn: Connection, table_name: str):
    """返回一个给定表的所有列名的列表"""
    # 执行SQL查询，获取表的所有列的信息
    columns = conn.execute(f"PRAGMA table_info('{table_name}');").fetchall()
    # 遍历columns,返回列名列表
    return [col[1] for col in columns]


def get_database_schema(conn: Connection) -> str:
    """获取数据库的Schema信息,包括表名及其包含的列表"""
    table_dicts = []  # 创建一个空的字典列表
    # 遍历数据库中的所有表
    for table_name in get_table_names(conn):
        columns_names = get_column_names(conn, table_name)  # 获取当前表的所有列名
        # 将表名和列名信息作为一个字典添加到列表中
        table_dicts.append({"table_name": table_name, "column_names": columns_names})

    # 格式化字典,转换成字符串返回
    database_schema_string = "\n".join(
        [
            f"Table: {table['table_name']}\nColumns: {', '.join(table['column_names'])}"
            for table in table_dicts
        ]
    )
    return database_schema_string


def generate_sql(db_schema: str, query: str) -> str:
    """调用LLM,利用工具调用能力,生成SQL语句"""

    dotenv.load_dotenv()
    client = OpenAI(base_url=os.getenv("OPENAI_API_BASE"))

    # 定义工具
    tool_desc = f"""SQL query extracting info to answer the user's question.
                SQL should be written using this database schema:
                {db_schema}
                The query should be returned in plain text, not in JSON.
                """
    tools = [
        {
            "type": "function",  # 工具类型为function函数调用
            "function": {  # 函数定义
                "name": "generate_sql",  # 函数名称
                "description": "Use this function to answer user questions about music. "
                               "Output should be a fully formed SQL query.",
                # 函数描述
                "parameters": {  # 函数参数定义
                    "type": "object",
                    "properties": {
                        "sql": {  # 参数名称
                            "type": "string",  # 参数类型
                            "description": tool_desc,  # 参数描述
                        },
                    },
                    "required": ["sql"],  # 必需的参数
                },
            }
        }
    ]

    # 创建消息列表
    messages = [
        {"role": "system",
         "content": "Answer user questions by generating SQL queries against the Chinook Music Database."},
        {"role": "user", "content": f"{query}"},
    ]

    # 执行工具调用,获取结果
    completion = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        tools=tools,
        tool_choice="auto"  # 工具选择模式为auto,表示由LLM自行推理,觉得是生成普通消息还是进行工具调用
    )

    # 将工具调用结果解析成sql字符串,并返回
    return json.loads(completion.choices[0].message.tool_calls[0].function.arguments).get("sql")


def exec_sql(conn: Connection, sql: str) -> str:
    """执行SQL,获取结果"""
    try:
        # 执行SQL,并将结果转换为字符串
        results = str(conn.execute(sql).fetchall())
    except Exception as e:
        # 执行失败,捕获异常并返回错误信息
        results = f"query failed with error: {e}"

    # 返回查询结果
    return results


if __name__ == '__main__':
    # 创建sqlite数据库连接
    conn = connect_db()

    # 获取数据库的schema信息
    db_schema = get_database_schema(conn)

    # 调用LLM,生成sql
    query = "What is the name of the album with the most tracks, and how many tracks does it have?"
    sql = generate_sql(db_schema, query)
    print(f"generate sql: \n{sql}")

    # 执行SQL,获取结果
    answer = exec_sql(conn, sql)
    print("------------------------------------")
    print(f"Q: {query}\nA: {answer}")
